package com.yoki.cloud.common.tenant;

import com.alibaba.druid.sql.ast.SQLExpr;
import com.alibaba.druid.sql.ast.SQLStatement;
import com.alibaba.druid.sql.ast.expr.SQLBinaryOpExpr;
import com.alibaba.druid.sql.ast.expr.SQLBinaryOperator;
import com.alibaba.druid.sql.ast.statement.*;
import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlCreateTableStatement;
import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlDeleteStatement;
import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlInsertStatement;
import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlUpdateStatement;
import com.alibaba.druid.sql.dialect.mysql.parser.MySqlStatementParser;
import com.alibaba.druid.sql.parser.SQLExprParser;
import com.alibaba.druid.sql.parser.SQLParserUtils;
import com.alibaba.druid.sql.visitor.SchemaStatVisitor;
import com.alibaba.druid.stat.TableStat;
import com.alibaba.druid.util.JdbcUtils;
import com.yoki.cloud.common.core.CoreRestStatus;
import com.yoki.cloud.common.core.CurrentUser;
import com.yoki.cloud.common.model.exception.BizException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.List;
import java.util.Map;
import java.util.Objects;

public class TenantSQLUtil {
    private static final Logger logger = LoggerFactory.getLogger(TenantSQLUtil.class);


    public static String processSQL(String sql) throws Exception { //使用mysql解析
        MySqlStatementParser sqlStatementParser = new MySqlStatementParser(sql);
        //解析select查询
        SQLStatement statement = sqlStatementParser.parseStatement();
        // 多租户只处理查询和添加
        if (statement instanceof SQLSelectStatement) {
            sql = processSelectSQL(statement);
        } else if (statement instanceof MySqlInsertStatement) {
            sql = processInsertSQL((MySqlInsertStatement) statement);
        } else if (statement instanceof MySqlDeleteStatement) {
            MySqlDeleteStatement ds = (MySqlDeleteStatement) statement;
            return processDeleteSQL(ds);
        } else if (statement instanceof MySqlCreateTableStatement) {

        } else if (statement instanceof MySqlUpdateStatement) {
            sql = processUpdateSQL((MySqlUpdateStatement) statement);
        } else if (statement instanceof SQLAlterTableStatement) {

        }
        return sql;
    }

    private static String processSelectSQL(SQLStatement stmt) {
//        List<Map<String, Object>> result = new ArrayList<>();
//        // SQLParserUtils.createSQLStatementParser可以将sql装载到Parser里面
//        SQLStatementParser parser = SQLParserUtils.createSQLStatementParser(sql, JdbcUtils.MYSQL);
//        // parseStatementList的返回值SQLStatement本身就是druid里面的语法树对象
//        List<SQLStatement> stmtList = parser.parseStatementList();
//        SQLStatement stmt = stmtList.get(0);

        SQLSelectStatement selectStmt = (SQLSelectStatement) stmt;
        // 拿到SQLSelect 通过在这里打断点看对象我们可以看出这是一个树的结构
        SQLSelect sqlselect = selectStmt.getSelect();
        SQLSelectQueryBlock query = (SQLSelectQueryBlock) sqlselect.getQuery();

        SchemaStatVisitor visitor = new SchemaStatVisitor();
        SQLTableSource sqlTableSource = query.getFrom();
        if (sqlTableSource == null) {// 没有查询表的SQL
            return stmt.toString();
        }
        sqlTableSource.accept(visitor);

        // 拼装租户查询条件，因为可能是联表查询，所以所有表都需要该条件，如果设置了别名需要改为别名
        StringBuffer tenantCondition = new StringBuffer();
        Map<String, String> aliasMap = visitor.getAliasMap();
        for (TableStat.Name tableName : visitor.getTables().keySet()) {
            // 查找表的别名
            String alias = null;
            for (Map.Entry<String, String> alias2Table : aliasMap.entrySet()) {
                String k = alias2Table.getKey(), v = alias2Table.getValue();
                // 这里别名map里包含有k、v都是原表名的数据
                if (!Objects.equals(k, v) && tableName.getName().equals(alias2Table.getValue())) {
                    alias = k;
                    break;
                }
            }

            if (tenantCondition.length() != 0) {
                tenantCondition.append(" AND ");
            }
            tenantCondition.append(alias != null ? alias : tableName).append(".").append(TENANT_FLAG_COLUMN).append(" = ").append(getSQLTenantFlag());
        }

        SQLExprParser constraintsParser = SQLParserUtils.createExprParser(tenantCondition.toString(), JdbcUtils.MYSQL);
        SQLExpr constraintsExpr = constraintsParser.expr();
        SQLExpr whereExpr = query.getWhere();

        // 修改where表达式
        if (whereExpr == null) {
            query.setWhere(constraintsExpr);
        } else {
            SQLBinaryOpExpr newWhereExpr = new SQLBinaryOpExpr(whereExpr, SQLBinaryOperator.BooleanAnd, constraintsExpr);
            query.setWhere(newWhereExpr);
        }
        sqlselect.setQuery(query);

        return sqlselect.toString();
    }


    private static final String TENANT_FLAG_COLUMN = "fd_reserved2";

    private static String processInsertSQL(MySqlInsertStatement insert) throws Exception {
        String tableName = removeBackquote(insert.getTableName().getSimpleName());
        // TODO simin
//            if(!isGlobalTable(tableName))
//                return sql;
//            if(!isInnerColExist(tableName))
//                return sql;

        if (insert.getQuery() != null)   // insert into tab select
            return null;

        StringBuilder sb = new StringBuilder(200)       // 指定初始容量可以提高性能
                .append("insert into ").append(tableName);

        List<SQLExpr> columns = insert.getColumns();

        int idx = -1;
        int colSize = -1;

        if (columns == null || columns.size() <= 0) { // insert 没有带列名：insert into t values(xxx,xxx)
            throw new RuntimeException("不支持不带列名的SQL");// TODO simin
        } else {  // insert 语句带有 列名
            sb.append("(");
            for (int i = 0; i < columns.size(); i++) {
                if (i < columns.size() - 1)
                    sb.append(columns.get(i).toString()).append(",");
                else
                    sb.append(columns.get(i).toString());
                String column = removeBackquote(insert.getColumns().get(i).toString());
                if (column.equalsIgnoreCase(TENANT_FLAG_COLUMN))
                    idx = i;
            }
            if (idx <= -1)
                sb.append(",").append(TENANT_FLAG_COLUMN);
            sb.append(")");
            colSize = columns.size();
        }

        sb.append(" values");
        List<SQLInsertStatement.ValuesClause> vcl = insert.getValuesList();
        if (vcl != null && vcl.size() > 1) {      // 批量insert
            for (int j = 0; j < vcl.size(); j++) {
                if (j != vcl.size() - 1)
                    appendValues(vcl.get(j).getValues(), sb, idx, colSize).append(",");
                else
                    appendValues(vcl.get(j).getValues(), sb, idx, colSize);
            }
        } else {  // 非批量 insert
            List<SQLExpr> valuse = insert.getValues().getValues();
            appendValues(valuse, sb, idx, colSize);
        }

        List<SQLExpr> dku = insert.getDuplicateKeyUpdate();
        if (dku != null && dku.size() > 0) {
            sb.append(" on duplicate key update ");
            for (int i = 0; i < dku.size(); i++) {
                SQLExpr exp = dku.get(i);
                if (exp != null) {
                    if (i < dku.size() - 1)
                        sb.append(exp.toString()).append(",");
                    else
                        sb.append(exp.toString());
                }
            }
        }
        return sb.toString();
    }

    private static String processDeleteSQL(MySqlDeleteStatement delete) {
        SQLTableSource ts = delete.getTableSource();
        if (ts != null && ts.toString().contains(",")) {
            // TODO
            throw new RuntimeException(ts.toString() + "Do not support Multiple-table delete syntax...");
        }

        String tableName = removeBackquote(delete.getTableName().getSimpleName());

        SQLExpr whereExpr = delete.getWhere();

        String tenantCondition = tableName + "." + TENANT_FLAG_COLUMN + "=" + getSQLTenantFlag();
        SQLExprParser constraintsParser = SQLParserUtils.createExprParser(tenantCondition, JdbcUtils.MYSQL);
        SQLExpr constraintsExpr = constraintsParser.expr();

        // 修改where表达式
        if (whereExpr == null) {
            delete.setWhere(constraintsExpr);
        } else {
            SQLBinaryOpExpr newWhereExpr = new SQLBinaryOpExpr(whereExpr, SQLBinaryOperator.BooleanAnd, constraintsExpr);
            delete.setWhere(newWhereExpr);
        }
        return delete.toString();
    }

    private static String processUpdateSQL(MySqlUpdateStatement update) {
        SQLTableSource ts = update.getTableSource();
        if (ts != null && ts.toString().contains(",")) {
            // TODO
            throw new RuntimeException(ts.toString() + "Do not support Multiple-table update syntax...");
        }

        String tableName = removeBackquote(update.getTableName().getSimpleName());

        SQLExpr whereExpr = update.getWhere();

        String tenantCondition = tableName + "." + TENANT_FLAG_COLUMN + "=" + getSQLTenantFlag();
        SQLExprParser constraintsParser = SQLParserUtils.createExprParser(tenantCondition, JdbcUtils.MYSQL);
        SQLExpr constraintsExpr = constraintsParser.expr();

        // 修改where表达式
        if (whereExpr == null) {
            update.setWhere(constraintsExpr);
        } else {
            SQLBinaryOpExpr newWhereExpr = new SQLBinaryOpExpr(whereExpr, SQLBinaryOperator.BooleanAnd, constraintsExpr);
            update.setWhere(newWhereExpr);
        }
        return update.toString();
    }

    private static StringBuilder appendValues(List<SQLExpr> valuse, StringBuilder sb, int idx, int colSize) {
        int size = valuse.size();
        if (size < colSize)
            size = colSize;

        sb.append("(");
        for (int i = 0; i < size; i++) {
            if (i < size - 1) {
                if (i != idx)
                    sb.append(valuse.get(i).toString()).append(",");
                else
                    sb.append(getSQLTenantFlag()).append(",");
            } else {
                if (i != idx) {
                    sb.append(valuse.get(i).toString());
                } else {
                    sb.append(getSQLTenantFlag());
                }
            }
        }
        if (idx <= -1)
            sb.append(",").append(getSQLTenantFlag());
        return sb.append(")");
    }

    public static String getTenantFlag() {
        String f = CurrentUser.getTenantFlag();
        if (f == null || f.length() == 0) {
            logger.error("未获取到租户标识");
            throw new BizException(CoreRestStatus.TENANT_NOT_FOUND);
        }
        return f;
    }

    public static String getSQLTenantFlag() {
        String f = CurrentUser.getTenantFlag();
        if (f == null || f.length() == 0) {
            throw new BizException(CoreRestStatus.TENANT_NOT_FOUND);
        }
        return "'" + getTenantFlag() + "'";
    }

    private static String removeBackquote(String s) {
        return s.replaceAll("`", "");
    }
}
